from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

import datetime  # For datetime objects
import os.path  # To manage paths
import sys  
sys.path.append('./backtrader')
import pandas as pd
from bk_manu_mark_data import manu_mark_process_stock_data


# Import the backtrader platform
import backtrader as bt

feture_output_file = 'stock_features_data.csv' 
scale_set = 1000
# Create a Stratey
class TestStrategy(bt.Strategy):
    params = (
        ('maperiod_5', 5),
        ('maperiod_10', 10),
        ('maperiod_20', 20),
        ('maperiod_30', 30),
        ('maperiod_60', 60),
        ('maperiod_120', 120),
        # 添加缺失的参数
        ('divergence_period', 30),
        ('top_divergence_period', 30),
        ('sell_ma_period', 5),
        ('bollinger_period', 20),
        ('bollinger_dev', 2.0),
        ('scale_set', 1000),
    )

    def log(self, txt, dt=None):
        ''' Logging function fot this strategy'''
        dt = dt or self.datas[0].datetime.date(0)
        print('%s, %s' % (dt.isoformat(), txt))

    def __init__(self):
        # Keep a reference to the "close" line in the data[0] dataseries
        self.dataclose = self.datas[0].close
        # 引用成交量数据
        self.volume = self.datas[0].volume
       
        # To keep track of pending orders and buy price/commission
        self.order = None
        self.buyprice = None
        self.buycomm = None
        self.trade_count = 0
        
        # 初始化数据存储列表，用于保存到CSV
        self.features_data = []


        # Add a MovingAverageSimple indicator
        # Close价格的SMA
        self.sma_close_5 = bt.indicators.SimpleMovingAverage(
            self.datas[0].close, period=self.params.maperiod_5)
        self.sma_close_10 = bt.indicators.SimpleMovingAverage(
            self.datas[0].close, period=self.params.maperiod_10)
        self.sma_close_20 = bt.indicators.SimpleMovingAverage(
            self.datas[0].close, period=self.params.maperiod_20)
        self.sma_close_30 = bt.indicators.SimpleMovingAverage(
            self.datas[0].close, period=self.params.maperiod_30)
        # 添加60日均线
        self.sma_close_60 = bt.indicators.SimpleMovingAverage(
            self.datas[0].close, period=self.params.maperiod_60)
        # 添加120日均线
        self.sma_close_120 = bt.indicators.SimpleMovingAverage(
            self.datas[0].close, period=self.params.maperiod_120)
        
        # Open价格的SMA
        self.sma_open_5 = bt.indicators.SimpleMovingAverage(
            self.datas[0].open, period=self.params.maperiod_5)
        self.sma_open_10 = bt.indicators.SimpleMovingAverage(
            self.datas[0].open, period=self.params.maperiod_10)
        self.sma_open_20 = bt.indicators.SimpleMovingAverage(
            self.datas[0].open, period=self.params.maperiod_20)
        self.sma_open_30 = bt.indicators.SimpleMovingAverage(
            self.datas[0].open, period=self.params.maperiod_30)
        # 添加60日均线
        self.sma_open_60 = bt.indicators.SimpleMovingAverage(
            self.datas[0].open, period=self.params.maperiod_60)
        self.sma_open_120 = bt.indicators.SimpleMovingAverage(
            self.datas[0].open, period=self.params.maperiod_120)
        
        # High价格的SMA
        self.sma_high_5 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high, period=self.params.maperiod_5)
        self.sma_high_10 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high, period=self.params.maperiod_10)
        self.sma_high_20 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high, period=self.params.maperiod_20)
        self.sma_high_30 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high, period=self.params.maperiod_30)
        # 添加60日均线
        self.sma_high_60 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high, period=self.params.maperiod_60)
        self.sma_high_120 = bt.indicators.SimpleMovingAverage(
            self.datas[0].high, period=self.params.maperiod_120)
        
        # Low价格的SMA
        self.sma_low_5 = bt.indicators.SimpleMovingAverage(
            self.datas[0].low, period=self.params.maperiod_5)
        self.sma_low_10 = bt.indicators.SimpleMovingAverage(
            self.datas[0].low, period=self.params.maperiod_10)
        self.sma_low_20 = bt.indicators.SimpleMovingAverage(
            self.datas[0].low, period=self.params.maperiod_20)
        self.sma_low_30 = bt.indicators.SimpleMovingAverage(
            self.datas[0].low, period=self.params.maperiod_30)
        # 添加60日均线
        self.sma_low_60 = bt.indicators.SimpleMovingAverage(
            self.datas[0].low, period=self.params.maperiod_60)
        self.sma_low_120 = bt.indicators.SimpleMovingAverage(
            self.datas[0].low, period=self.params.maperiod_120)
        
        # Volume成交量的SMA
        self.sma_volume_5 = bt.indicators.SimpleMovingAverage(
            self.datas[0].volume, period=self.params.maperiod_5)
        self.sma_volume_10 = bt.indicators.SimpleMovingAverage(
            self.datas[0].volume, period=self.params.maperiod_10)
        self.sma_volume_20 = bt.indicators.SimpleMovingAverage(
            self.datas[0].volume, period=self.params.maperiod_20)
        self.sma_volume_30 = bt.indicators.SimpleMovingAverage(
            self.datas[0].volume, period=self.params.maperiod_30)
        # 添加60日均线
        self.sma_volume_60 = bt.indicators.SimpleMovingAverage(
            self.datas[0].volume, period=self.params.maperiod_60)
        self.sma_volume_120 = bt.indicators.SimpleMovingAverage(
            self.datas[0].volume, period=self.params.maperiod_120)
        
        # 为了保持向后兼容，保留原来的sma引用
        self.sma = self.sma_close_5
        self.sma10 = self.sma_close_10
        self.sma20 = self.sma_close_20
        self.sma30 = self.sma_close_30
        self.sma120 = self.sma_close_120
        
        # 添加5日线上穿120日线的交叉指标
        self.crossover_5_120 = bt.ind.CrossOver(self.sma, self.sma120)

        
        # Indicators for the plotting show
        # bt.indicators.ExponentialMovingAverage(self.datas[0], period=25)
        # bt.indicators.WeightedMovingAverage(self.datas[0], period=25,
        #                                     subplot=True)
        # bt.indicators.StochasticSlow(self.datas[0])
        self.macd_histo = bt.indicators.MACDHisto(self.datas[0])
        # add diff and dea indicator
        self.macd = bt.indicators.MACD(self.datas[0])
        # MACD包含了三个值：macd (DIFF线), signal (DEA线), histo (柱状图)
        self.diff = self.macd.macd  # DIFF线 (快速EMA - 慢速EMA)
        self.dea = self.macd.signal  # DEA线 (DIFF的EMA平滑线)
        
        # 添加最低价指标用于底背离判断
        self.lowest_period = bt.indicators.Lowest(self.datas[0].low, period=self.params.divergence_period)
        
        # 添加最高价指标用于顶背离判断
        self.highest_period = bt.indicators.Highest(self.datas[0].high, period=self.params.top_divergence_period)
        
        # 添加卖出信号用的均线指标
        self.sell_ma = bt.indicators.SimpleMovingAverage(
            self.datas[0], period=self.params.sell_ma_period)
        
        # self.macd = bt.indicators.MACD(self.datas[0])
        # rsi = bt.indicators.RSI(self.datas[0])
        # bt.indicators.SmoothedMovingAverage(rsi, period=10)
        # bt.indicators.ATR(self.datas[0], plot=False)
        
        # 添加布林带指标
        self.bollinger = bt.indicators.BollingerBands(
            self.datas[0],
            period=self.params.bollinger_period,
            devfactor=self.params.bollinger_dev
        )

    def notify_order(self, order):
        if order.status in [order.Submitted, order.Accepted]:
            # Buy/Sell order submitted/accepted to/by broker - Nothing to do
            return

        # Check if an order has been completed
        # Attention: broker could reject order if not enough cash
        if order.status in [order.Completed]:
            if order.isbuy():
                # self.log(
                #     'BUY EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                #     (order.executed.price,
                #      order.executed.value,
                #      order.executed.comm))

                self.buyprice = order.executed.price
                self.buycomm = order.executed.comm
            else:  # Sell
                # self.log('SELL EXECUTED, Price: %.2f, Cost: %.2f, Comm %.2f' %
                #          (order.executed.price,
                #           order.executed.value,
                #           order.executed.comm))
                pass

            self.bar_executed = len(self)

        elif order.status in [order.Canceled, order.Margin, order.Rejected]:
            self.log('Order Canceled/Margin/Rejected')

        self.order = None

    def notify_trade(self, trade):
        if not trade.isclosed:
            return

        # 统计交易次数
        self.trade_count += 1
        
        # self.log('OPERATION PROFIT, GROSS %.2f, NET %.2f' %
        #          (trade.pnl, trade.pnlcomm))


    def start(self):
        # 记录开始信息
        self.start_cash = self.broker.getvalue()
        # 暂时将start_date设为None，稍后在next方法中设置
        self.start_date = None
        self.log('start  : %.2f' % self.start_cash)
        
    def stop(self):
        # 记录结束信息
        self.end_date = self.datas[0].datetime.date(0)
        final_cash = self.broker.getvalue()
        self.log('============>period={}'.format(self.params.bollinger_period))
        # self.log('end date: %s, last cash: %.2f' % (self.end_date, final_cash))计算统计数据
        total_return = ((final_cash - self.start_cash) / self.start_cash) * 100
        
        # 将收集的数据保存为CSV文件
        if self.features_data:
            df = pd.DataFrame(self.features_data)
            # 将所有数值型字段缩小1000倍并保留3位小数点
            numeric_columns = df.select_dtypes(include=['float64', 'int64']).columns
            for col in numeric_columns:
                if self.params.scale_set > 1 :
                    df[col] = (df[col] / 1000).round(3)
                else:
                    df[col] = (df[col]).round(3)

            df.to_csv(feture_output_file, index=False)
            self.log(f'特征数据已保存到----> {feture_output_file}')
        
    
    def next(self):
        # 记录第一个交易日作为开始日期
        if self.start_date is None:
            self.start_date = self.datas[0].datetime.date(0)
            self.log('start date: %s' % self.start_date)
        
        # 获取当前日期
        current_date = self.datas[0].datetime.date(0)
        
        # 获取价格数据
        open_price = self.datas[0].open[0]
        high_price = self.datas[0].high[0]
        low_price = self.datas[0].low[0]
        close_price = self.datas[0].close[0]
        volume = self.datas[0].volume[0]
        
        # 计算当日close均值（这里使用当日close作为均值，如果需要其他计算方式可以修改）
        close_mean = close_price
        
        # 收集数据到字典
        data_row = {
            'Date': current_date,
            'Open': open_price,
            'High': high_price,
            'Low': low_price,
            'Close': close_price,
            'Volume': volume,
            'Close_Mean': close_mean,
            # Close价格的SMA指标
            'Close_SMA5': self.sma_close_5[0],
            'Close_SMA10': self.sma_close_10[0],
            'Close_SMA20': self.sma_close_20[0],
            'Close_SMA30': self.sma_close_30[0],
            'Close_SMA60': self.sma_close_60[0],
            'Close_SMA120': self.sma_close_120[0],
            # Open价格的SMA指标
            'Open_SMA5': self.sma_open_5[0],
            'Open_SMA10': self.sma_open_10[0],
            'Open_SMA20': self.sma_open_20[0],
            'Open_SMA30': self.sma_open_30[0],
            'Open_SMA60': self.sma_open_60[0],
            'Open_SMA120': self.sma_open_120[0],
            # High价格的SMA指标
            'High_SMA5': self.sma_high_5[0],
            'High_SMA10': self.sma_high_10[0],
            'High_SMA20': self.sma_high_20[0],
            'High_SMA30': self.sma_high_30[0],
            'High_SMA60': self.sma_high_60[0],
            'High_SMA120': self.sma_high_120[0],
            # Low价格的SMA指标
            'Low_SMA5': self.sma_low_5[0],
            'Low_SMA10': self.sma_low_10[0],
            'Low_SMA20': self.sma_low_20[0],
            'Low_SMA30': self.sma_low_30[0],
            'Low_SMA60': self.sma_low_60[0],
            'Low_SMA120': self.sma_low_120[0],
            # Volume成交量的SMA指标
            'Volume_SMA5': self.sma_volume_5[0],
            'Volume_SMA10': self.sma_volume_10[0],
            'Volume_SMA20': self.sma_volume_20[0],
            'Volume_SMA30': self.sma_volume_30[0],
            'Volume_SMA60': self.sma_volume_60[0],
            'Volume_SMA120': self.sma_volume_120[0]
        }
        
        # 将数据添加到列表
        self.features_data.append(data_row)
        
        pass
        
        
        
        
data_path_global = './csv_data/SH_000300_stock_data.csv'
def   bk_create_feture_data(data_file_path,  date_start ,date_end, scale_set=1):
    print(f'bk_create_feture_data: {date_start} ~ {date_end}')

    # Create a cerebro entity
    cerebro = bt.Cerebro()

    # Add a strategy with parameter optimization
    # strats = cerebro.optstrategy(TestStrategy, bollinger_period=range(20, 90, 5))
    
    # # Add a strategy
    cerebro.addstrategy(TestStrategy, scale_set=scale_set)

    # Datas are in a subfolder of the samples. Need to find where the script is
    # because it could have been called from anywhere
    modpath = os.path.dirname(os.path.abspath(sys.argv[0]))
    datapath = os.path.join(modpath, data_file_path)

    # Create a Data Feed
    data = bt.feeds.YahooFinanceCSVData(
        dataname=datapath,
        # Do not pass values before this date
        # fromdate=datetime.datetime(2023, 8, 3),
        fromdate=date_start,
        # Do not pass values after this date
        # todate=datetime.datetime(2025, 9, 5),
        todate=date_end,
        # Do not pass values after this date
        reverse=False)

    # Add the Data Feed to Cerebro
    cerebro.adddata(data)

    # Set our desired cash start
    cerebro.broker.setcash(10000000.0)

    # Add a FixedSize sizer according to the stake
    cerebro.addsizer(bt.sizers.FixedSize, stake=10)

    # Set the commission
    cerebro.broker.setcommission(commission=0.05)

    # Print out the starting conditions
    print('Starting Portfolio Value: %.2f' % cerebro.broker.getvalue())

    # Run over everything
    cerebro.run()

    # Print out the final result
    final_value = cerebro.broker.getvalue()
    print('Final Portfolio Value: %.2f' % final_value)
    out_file_str = 'feture_data_%s_%s.csv' % (date_start.strftime('%Y%m%d'), date_end.strftime('%Y%m%d'))
    result_df = manu_mark_process_stock_data(feture_output_file, out_file_str, True)
    return result_df, out_file_str
    pass
if __name__ == '__main__':
    bk_create_feture_data(data_path_global, 
        date_start=datetime.datetime(2010, 1, 4),
        date_end=datetime.datetime(2020, 1, 4), scale_set=1000)

    
    #============================================================================
    # bk_create_feture_data(data_path_global, 
    #     date_start=datetime.datetime(2023, 8, 3),
    #     date_end=datetime.datetime(2025, 9, 5))

    pass
